# Copyright 2024 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""Internal implementation of proto compilation.

# Overview of implementation

(If you just want to use the macros, see their docstrings; this section is
intended to orient future maintainers.)

Proto code generation is carried out by the pwpb_proto_library,
nanopb_proto_library, pw_raw_rpc_proto_library and pw_nanopb_rpc_proto_library
rules using aspects
(https://docs.bazel.build/versions/main/skylark/aspects.html).

As an example, pwpb_proto_library has a single proto_library as a dependency,
but that proto_library may depend on other proto_library targets; as a result,
the generated .pwpb.h file #include's .pwpb.h files generated from the
dependency proto_libraries. The aspect propagates along the proto_library
dependency graph, running the proto compiler on each proto_library in the
original target's transitive dependencies, ensuring that we're not missing any
.pwpb.h files at C++ compile time.

Although we have a separate rule for each protocol compiler plugin
(pwpb_proto_library, nanopb_proto_library, pw_raw_rpc_proto_library,
pw_nanopb_rpc_proto_library), they actually share an implementation
(compile_proto) and use similar aspects, all generated by
proto_compiler_aspect.
"""

load("@bazel_skylib//lib:paths.bzl", "paths")
load("@com_google_protobuf//bazel/common:proto_info.bzl", "ProtoInfo")
load(
    "//pw_build/bazel_internal:pigweed_internal.bzl",
    _compile_cc = "compile_cc",
)
load("//pw_protobuf_compiler:pw_proto_filegroup.bzl", "PwProtoOptionsInfo")

PwProtoInfo = provider(
    "Returned by PW proto compilation aspect",
    fields = {
        "hdrs": "generated C++ header files",
        "includes": "include paths for generated C++ header files",
        "srcs": "generated C++ src files",
    },
)

def compile_proto(ctx):
    """Implementation of the proto codegen rule.

    The work of actually generating the code is done by the aspect, so here we
    compile and return a CcInfo to link against.

    Args:
      ctx: Rule context object (https://bazel.build/rules/lib/builtins/ctx).

    Returns:
      A CcInfo provider.
    """

    # Note that we don't distinguish between the files generated from the
    # target, and the files generated from its dependencies. We return all of
    # them together, and in pw_proto_library expose all of them as hdrs.
    # Pigweed's plugins happen to only generate .h files, so this works, but
    # strictly speaking we should expose only the files generated from the
    # target itself in hdrs, and place the headers generated from dependencies
    # in srcs. We don't perform layering_check in Pigweed, so this is not a big
    # deal.
    #
    # TODO: b/234873954 - Tidy this up.
    all_srcs = []
    all_hdrs = []
    all_includes = []
    for dep in ctx.attr.protos:
        for f in dep[PwProtoInfo].hdrs:
            all_hdrs.append(f)
        for f in dep[PwProtoInfo].srcs:
            all_srcs.append(f)
        for i in dep[PwProtoInfo].includes:
            all_includes.append(i)

    return _compile_cc(
        ctx,
        all_srcs,
        all_hdrs,
        ctx.attr.deps,
        all_includes,
        defines = [],
    )

def _options_symlink_path(options_file, workspace_root, proto_source_root, import_prefix, strip_import_prefix):
    path_in_module = paths.relativize(options_file.path, workspace_root)

    if strip_import_prefix:
        stripped_path = paths.relativize(path_in_module, strip_import_prefix.lstrip("/"))
    else:
        stripped_path = path_in_module

    if import_prefix:
        extended_path = paths.join(import_prefix, stripped_path)
    else:
        extended_path = stripped_path

    return paths.join(proto_source_root, extended_path)

def _proto_compiler_aspect_impl(target, ctx):
    for excluded_target in ctx.attr._excluded_targets:
        if excluded_target.label == target.label:
            return PwProtoInfo(srcs = [], hdrs = [], includes = [])

    # List the files we will generate for this proto_library target.
    proto_info = target[ProtoInfo]

    srcs = []
    hdrs = []

    # Setup the output root for the plugin to point to targets output
    # directory. This allows us to declare the location of the files that protoc
    # will output in a way that `ctx.actions.declare_file` will understand,
    # since it works relative to the target.
    out_path = ctx.bin_dir.path
    if target.label.workspace_root:
        out_path += "/" + target.label.workspace_root
    if target.label.package:
        out_path += "/" + target.label.package

    # Add location of headers to cc include path.
    # Depending on prefix rules, the include path can be directly from the
    # output path, or underneath the package.
    includes = [out_path]

    for src in proto_info.direct_sources:
        # Get the relative import path for this .proto file.
        src_rel = paths.relativize(src.path, proto_info.proto_source_root)
        proto_dir = paths.dirname(src_rel)

        # Add location of headers to cc include path.
        includes.append("{}/{}".format(out_path, src.owner.package))

        for ext in ctx.attr._extensions:
            # Declare all output files, in target package dir.
            generated_filename = src.basename[:-len("proto")] + ext
            if proto_dir:
                out_file_name = "{}/{}".format(
                    proto_dir,
                    generated_filename,
                )
            else:
                out_file_name = generated_filename

            out_file = ctx.actions.declare_file(out_file_name)

            if ext.endswith(".h"):
                hdrs.append(out_file)
            else:
                srcs.append(out_file)

    # The proto_source_root may be prefixed with the output directory. But it
    # may not. Ensure that there is no such prefix, which is intended to become
    # the only case one day. See
    # https://github.com/protocolbuffers/protobuf/blob/069a66850d1d8bb83c1ca1eb5bdee87525290584/bazel/private/proto_info.bzl#L154-L165
    relative_proto_source_root = proto_info.proto_source_root
    if relative_proto_source_root.startswith(out_path):
        relative_proto_source_root = paths.relativize(relative_proto_source_root, out_path)

    # Symlink the .options files into the proto_source_root, so that they can be
    # found by protoc plugins regardless of [strip_]import_prefix attribute
    # values.
    #
    # For example, say we have a proto_library in //a/b/BUILD.bazel with
    # strip_import_prefix = b and import_prefix = xyz. Then, the `.proto` files
    # will live in the directory,
    #
    # bazel-bin/a/b/_virtual_imports/a/xyz/
    #
    # What we do here is move the `.options` files to the same directory. Later
    # on, we'll provide `bazel-bin/a/b/_virtual_imports` to the protoc plugin's
    # search path via `--custom_opt=-I`. This way, the proto and options files
    # will be alongside each other, as the plugins expect.
    symlinks = []
    for src in ctx.rule.attr.srcs:
        if PwProtoOptionsInfo in src:
            for options_file in src[PwProtoOptionsInfo].options_files.to_list():
                path_to_options_file = _options_symlink_path(
                    options_file,
                    target.label.workspace_root,
                    relative_proto_source_root,
                    ctx.rule.attr.import_prefix,
                    ctx.rule.attr.strip_import_prefix,
                )
                options_symlink_out = ctx.actions.declare_file(path_to_options_file)
                ctx.actions.symlink(output = options_symlink_out, target_file = options_file)
                symlinks.append(options_symlink_out)

    # List the `.options` files from any `pw_proto_filegroup` targets listed
    # under this target's `srcs`.
    options_files = [
        options_file
        for src in ctx.rule.attr.srcs
        if PwProtoOptionsInfo in src
        for options_file in src[PwProtoOptionsInfo].options_files.to_list()
    ]

    args = ctx.actions.args()
    for path in proto_info.transitive_proto_path.to_list():
        args.add("-I{}".format(path))

    args.add("--plugin=protoc-gen-custom={}".format(ctx.executable._protoc_plugin.path))
    args.add("--custom_opt=-I{}".format(paths.join(out_path, relative_proto_source_root)))

    for plugin_option in ctx.attr._plugin_options:
        # If the plugin supports directly specifying the location of the options files, pass them here.
        if plugin_option == "--options-file={}":
            for options_file in options_files:
                plugin_options_arg = plugin_option.format(options_file.path)
                args.add("--custom_opt={}".format(plugin_options_arg))
            continue
        args.add("--custom_opt={}".format(plugin_option))

    args.add("--custom_out={}".format(out_path))
    args.add_all(proto_info.direct_sources)

    all_tools = [
        ctx.executable._protoc,
        ctx.executable._protoc_plugin,
    ]

    ctx.actions.run(
        inputs = depset(
            direct = proto_info.direct_sources +
                     proto_info.transitive_sources.to_list() +
                     options_files + symlinks,
            transitive = [proto_info.transitive_descriptor_sets],
        ),
        progress_message = "Generating %s C++ files for %s" % (ctx.attr._extensions, ctx.label.name),
        tools = all_tools,
        outputs = srcs + hdrs,
        executable = ctx.executable._protoc,
        arguments = [args],
        env = {
            # This effectively pre-adopts
            # https://github.com/nanopb/nanopb/pull/1038, silencing an annoying
            # warning in nanopb 0.4.9.1.
            "NANOPB_PB2_NO_REBUILD": "1",
            # The nanopb protobuf plugin likes to compile some temporary protos
            # next to source files. This forces them to be written to Bazel's
            # genfiles directory.
            "NANOPB_PB2_TEMP_DIR": str(ctx.genfiles_dir),
        },
    )

    transitive_srcs = srcs
    transitive_hdrs = hdrs
    transitive_includes = includes
    for dep in ctx.rule.attr.deps:
        transitive_srcs += dep[PwProtoInfo].srcs
        transitive_hdrs += dep[PwProtoInfo].hdrs
        transitive_includes += dep[PwProtoInfo].includes
    return [PwProtoInfo(
        srcs = transitive_srcs,
        hdrs = transitive_hdrs,
        includes = transitive_includes,
    )]

def proto_compiler_aspect(extensions, protoc_plugin, plugin_options = [], excluded_targets = []):
    """Returns an aspect that runs the proto compiler.

    The aspect propagates through the deps of proto_library targets, running
    the proto compiler with the specified plugin for each of their source
    files. The proto compiler is assumed to produce one output file per input
    .proto file. That file is placed under bazel-bin at the same path as the
    input file, but with the specified extension (i.e., with _extensions = [
    .pwpb.h], the aspect converts pw_log/log.proto into
    bazel-bin/pw_log/log.pwpb.h).

    The aspect returns a provider exposing all the File objects generated from
    the dependency graph.
    """
    return aspect(
        attr_aspects = ["deps"],
        attrs = {
            "_excluded_targets": attr.label_list(
                default = excluded_targets,
                doc = "List of targets at which the aspect should stop propagating.",
            ),
            "_extensions": attr.string_list(default = extensions),
            "_plugin_options": attr.string_list(
                default = plugin_options,
                doc = "List of options to pass to the protoc plugin.",
            ),
            "_protoc": attr.label(
                default = Label("@com_google_protobuf//:protoc"),
                executable = True,
                doc = "Path to the protoc binary.",
                cfg = "exec",
            ),
            "_protoc_plugin": attr.label(
                default = Label(protoc_plugin),
                executable = True,
                doc = "Protoc plugin to invoke.",
                cfg = "exec",
            ),
        },
        implementation = _proto_compiler_aspect_impl,
        provides = [PwProtoInfo],
    )
